
package org.jdamico.socks.server.impl;

import java.io.IOException;
import java.io.InterruptedIOException;
import java.net.DatagramPacket;
import java.net.DatagramSocket;
import java.net.InetAddress;
import java.net.UnknownHostException;
import java.util.Map;

import org.jdamico.socks.server.commons.Configs;
import org.jdamico.socks.server.commons.DebugLog;
import org.jdamico.socks.server.commons.NetworkUtils;
import org.jdamico.socks.server.commons.UserManager;
import org.jdamico.socks.server.commons.Utils;


public class Socks5Impl extends Socks4Impl implements SocksCommonInterface {

	protected DatagramSocket mDGSocket = null;
	protected DatagramPacket mDGPack = null;

	private InetAddress mUDP_IA = null;
	private int mUDP_port = 0;

	public byte RSV_VAL; // 保留字
	public byte addressType; // 地址类型

	static final int ADDR_Size[] = { -1, // '00' No such AType
			4, // '01' IP v4 - ipv4 占4个字节
			-1, // '02' No such AType
			-1, // '03' 域名地址，地址位开始的第一个字节位域名长度
			16 // '04' IP v6- ipv6 占16个字节
	};

	public Socks5Impl(ProxyHandler Parent) {
		super(Parent);
		mDST_Addr = new byte[Configs.MAX_ADDR_LEN];
	}

	public InetAddress calcInetAddress(byte addressType, byte[] addr) {
		InetAddress IA = null;

		switch (addressType) {
		case 0x01:
			IA = Utils.calcInetAddress(addr);
			break;

		case 0x03:
			IA = getDomainAddress(addr);
			break;
		default:
			return null;
		}
		return IA;
	}

	private InetAddress getDomainAddress(byte[] addr) {

		InetAddress address = null;
		if (addr[0] <= 0) {
			DebugLog.getInstance().error("SOCKS 5 - calcInetAddress() : BAD IP in command - size : " + addr[0]);
			return null;
		}
		String sIA = "";
		for (int i = 1; i <= addr[0]; i++) {
			sIA += (char) addr[i];
		}
		try {
			address = InetAddress.getByName(sIA);
		} catch (UnknownHostException e) {
			return null;
		}
		return address;
	}

	public boolean calculateAddress() {

		mServer = calcInetAddress(addressType, mDST_Addr);
		mServerPort = Utils.calcPort(mDST_Port[0], mDST_Port[1]);
		mClientIP = mProxyHandler.mClientSocket.getInetAddress();
		mClientPort = mProxyHandler.mClientSocket.getPort();

		return ((mServer != null) && (mServerPort >= 0));
	}

	public void authenticate(byte SOCKS_Ver) throws Exception {
		super.authenticate(SOCKS_Ver);

		if (SOCKS_Version == Configs.SOCKS5_Version) {
			
			if(!Configs.SHOULD_AUTH_USER) {
				if (!checkNoneAuthentication()) {
					refuseAuthentication("SOCKS 5 - Not Supported Authentication!");
					throw new Exception("SOCKS 5 - Not Supported Authentication.");
				}
				acceptAuthentication();
			}else {
				if(!checkUserAuthentication()) {
					refuseAuthentication("SOCKS 5 - Not Supported Authentication!");
					throw new Exception("SOCKS 5 - Not Supported Authentication.");
				}
				needUserAuthentication();
				
				int rfc_Version = getByte();
				if(rfc_Version!=Configs.AUTH_VERSION_1) {
					refuseAuthentication("Socks 5 Error,Don't Support the auth protocol!");
					throw new Exception("Socks 5 Error,Don't Support the auth protocol!");
				}
				String userName = readUserName();
				String password = readUserPassword();
				
				boolean findUser = UserManager.getInstance().handleUserInfo(userName,password);
				if(!findUser) {
					refuseAuthentication("Socks 5 Error,Not found the user");
					throw new Exception("Socks 5 Error,Not found the user");
				}else {	
					DebugLog.getInstance().print("UserInfo is valided");
					acceptAuthentication();
					this.mUID = UserManager.getInstance().createToken(mProxyHandler.mClientSocket.getInetAddress(), mProxyHandler.mClientSocket.getLocalPort(), userName+"@"+password);
				} 
				
			}
			
		} else {
			refuseAuthentication("Incorrect SOCKS version : " + SOCKS_Version);
			throw new Exception("Not Supported SOCKS Version -'" + SOCKS_Version + "'");
		}
	}
	
	private String readUserName() throws IOException {
		byte usenameLength = (byte) getByte();
		byte[] buf = new byte[usenameLength];
		for (int i = 0; i < usenameLength; i++) {
			buf[i] =  (byte) getByte();
		}
		String userName = new String(buf, 0, usenameLength);
		return userName;
	}
	
	private String readUserPassword() throws IOException {
		byte passwordLength = (byte) getByte();
		byte[] buf = new byte[passwordLength];
		for (int i = 0; i < passwordLength; i++) {
			buf[i] =  (byte) getByte();
		}
		String password = new String(buf, 0, passwordLength);
		return password;
	}

	public void refuseAuthentication(String msg) {
		DebugLog.getInstance().println("SOCKS 5 - Refuse Authentication: '" + msg + "'");
		mProxyHandler.sendToClient(Configs.SRE_Refuse);
	}

	public void acceptAuthentication() {
		DebugLog.getInstance().println("SOCKS 5 - Accepts Auth. method 'NO_AUTH'");
		byte[] tSRE_Accept = Configs.SRE_Accept;
		mProxyHandler.sendToClient(tSRE_Accept);
	}
	
	public void needUserAuthentication() {
		DebugLog.getInstance().println("SOCKS 5 - Accepts Auth. method 'AUTH USER'");
		byte[] tSRE_Accept = Configs.SRE_Auth_User;
		mProxyHandler.sendToClient(tSRE_Accept);
	}

	public boolean checkNoneAuthentication() throws Exception {
		// boolean Have_NoAuthentication = false;
		int method = 0x00;
		return HasAuthMethod(method);
	}
	
	public boolean checkUserAuthentication() throws Exception {
		// boolean Have_NoAuthentication = true;
		int method = 0x02;
		return HasAuthMethod(method);
	}

	private boolean HasAuthMethod(int code) {
		byte Methods_Num = getByte();
		boolean hasMethod = false;
		for (int i = 0; i < Methods_Num; i++) {
			if(getByte()==code) {
				hasMethod = true;
			}
		}
		return hasMethod;
	}

	public void checkClientCommand() throws Exception {
		int Addr_Len;

		SOCKS_Version = getByte();
		mSocksCommand = getByte();
		RSV_VAL = getByte();
		addressType = getByte();

		Addr_Len = ADDR_Size[addressType];
		mDST_Addr[0] = getByte();
		if (addressType == 0x03) {
			Addr_Len = mDST_Addr[0] + 1;
		}

		for (int i = 1; i < Addr_Len; i++) {
			mDST_Addr[i] = getByte();
		}
		mDST_Port[0] = getByte();
		mDST_Port[1] = getByte();

		if (SOCKS_Version != Configs.SOCKS5_Version) {
			DebugLog.getInstance().println("SOCKS 5 - Incorrect SOCKS Version of Command: " + SOCKS_Version);
			refuseCommand((byte) 0xFF);
			throw new Exception("Incorrect SOCKS Version of Command: " + SOCKS_Version);
		}

		if ((mSocksCommand < Configs.SC_CONNECT) || (mSocksCommand > Configs.SC_UDP)) {
			DebugLog.getInstance().error("SOCKS 5 - GetClientCommand() - Unsupported Command : \""
					+ Utils.getCommandName(mSocksCommand) + "\"");
			refuseCommand((byte) 0x07);
			throw new Exception("SOCKS 5 - Unsupported Command: \"" + mSocksCommand + "\"");
		}

		if (addressType == 0x04) {
			DebugLog.getInstance().error("SOCKS 5 - checkClientCommand() - Unsupported Address Type - IP v6");
			refuseCommand((byte) 0x08);
			throw new Exception("Unsupported Address Type - IP v6");
		}

		if ((addressType >= 0x04) || (addressType <= 0)) {
			DebugLog.getInstance().error("SOCKS 5 - checkClientCommand() - Unsupported Address Type: " + addressType);
			refuseCommand((byte) 0x08);
			throw new Exception("SOCKS 5 - Unsupported Address Type: " + addressType);
		}

		if (!calculateAddress()) { // Gets the IP Address
			refuseCommand((byte) 0x04);// Host Not Exists...
			String hostInfo = mServer!=null ?mServer.toString() :"";
			throw new Exception("SOCKS 5 - Unknown Host/IP address '" + hostInfo + "'");
		}

		DebugLog.getInstance()
				.println("SOCKS 5 - Accepted SOCKS5 Command: \"" + Utils.getCommandName(mSocksCommand) + "\"");
	}

	/**
	 * 身份验证阶段用于回复客户端
	 */
	public void replyCommand(byte replyCode) {
		DebugLog.getInstance().println("SOCKS 5 - Proxy Reply to Client \"" + Utils.getReplyName(replyCode) + "\"");

		int port = 0;
		byte[] REPLY = new byte[10];
		byte IP[] = {0,0,0,0};

		if (mProxyHandler.mServerSocket != null) {
			port = mProxyHandler.mServerSocket.getLocalPort();
		} else {
			port = 0;
		}

		REPLY[0] = Configs.SOCKS5_Version;
		REPLY[1] = replyCode; // 响应码
		REPLY[2] = 0x00; // 保留位 '00'
		REPLY[3] = Configs.ADDRESS_IPV4; // IP类型
		//代理服务器返回BND.Addr和BND.Addr,这里设置位0.0.0.0的地址，没有实际意义，客户端也不会去校验
		REPLY[4] = IP[0];  
		REPLY[5] = IP[1];
		REPLY[6] = IP[2];
		REPLY[7] = IP[3];
		REPLY[8] = (byte) ((port & 0xFF00) >> 8);// Port High
		REPLY[9] = (byte) (port & 0x00FF); // Port Low

		mProxyHandler.sendToClient(REPLY);// BND.PORT
	}

	/**
	 * 地址绑定阶段
	 */
	public void bindReply(byte replyCode, InetAddress localhostAddress, int port) {
		byte IP[] = { 0, 0, 0, 0 };

		DebugLog.getInstance().println("BIND Reply to Client \"" + Utils.getReplyName(replyCode) + "\"");
		byte[] REPLY = new byte[10];
		if (localhostAddress != null){
			IP = localhostAddress.getAddress();
		}

		REPLY[0] = Configs.SOCKS5_Version;
		REPLY[1] = (byte) ((int) replyCode - 90); // Reply Code;
		REPLY[2] = 0x00; // Reserved '00'
		REPLY[3] = Configs.ADDRESS_IPV4; // IP ver.4 Type
		REPLY[4] = IP[0];
		REPLY[5] = IP[1];
		REPLY[6] = IP[2];
		REPLY[7] = IP[3];
		REPLY[8] = (byte) ((port & 0xFF00) >> 8);
		REPLY[9] = (byte) (port & 0x00FF);

		if (mProxyHandler.isActive()) {
			mProxyHandler.sendToClient(REPLY);
		} else {
			DebugLog.getInstance().println("BIND - Closed Client Connection");
		}
	}

	/**
	 * udp连接确认
	 * @param replyCode
	 * @param address
	 * @param port
	 * @throws IOException
	 */
	public void udpReply(byte replyCode, InetAddress address, int port) throws IOException {

		DebugLog.getInstance().println("Reply to Client \"" + Utils.getReplyName(replyCode) + "\"");
		if (mProxyHandler.mClientSocket == null) {
			DebugLog.getInstance().println("Error in UDP_Reply() - Client socket is NULL");
		}
		byte[] IP = address.getAddress();
		byte[] REPLY = new byte[10];
		REPLY[0] = Configs.SOCKS5_Version;
		REPLY[1] = replyCode; // 响应吗;
		REPLY[2] = 0x00; // 保留字节位 ，值位'00'
		REPLY[3] = 0x01; //地址类型
		REPLY[4] = IP[0];
		REPLY[5] = IP[1];
		REPLY[6] = IP[2];
		REPLY[7] = IP[3];

		REPLY[8] = (byte) ((port & 0xFF00) >> 8);// 小端存储，高位在前
		REPLY[9] = (byte) (port & 0x00FF); // 小端存储，低位在后

		mProxyHandler.sendToClient(REPLY);// BND.PORT
	}

	public void udp() throws IOException {

		try {
			mDGSocket = new DatagramSocket();
			initUdpInOut();
		} catch (IOException e) {
			refuseCommand((byte) 0x05); // Connection Refused
			throw new IOException("Connection Refused - FAILED TO INITIALIZE UDP Association.");
		}

		//InetAddress MyIP = mProxyHandler.mClientSocket.getLocalAddress();
		Map<String, Object> inetIdent = NetworkUtils.getLocalInetIdent();
		
		InetAddress proxyHost = InetAddress.getByName((String) inetIdent.get("ip"));
		int proxyPort = mDGSocket.getLocalPort();

		udpReply((byte) 0, proxyHost, proxyPort);
		DebugLog.getInstance().println("UDP Listen at: <" + proxyHost.toString() + ":" + proxyPort + ">");

		while (mProxyHandler.checkClientData() >= 0) {
			processUdp();
			Thread.yield();
		}
		DebugLog.getInstance().println("UDP - Closed TCP Master of UDP Association");
	}

	private void initUdpInOut() throws IOException {

		mDGSocket.setSoTimeout(Configs.DEFAULT_PROXY_TIMEOUT);
		mProxyHandler.mBuffer = new byte[Configs.DEFAULT_BUF_SIZE];
		mDGPack = new DatagramPacket(mProxyHandler.mBuffer, Configs.DEFAULT_BUF_SIZE);
	}

	private byte[] addDgpHead(byte[] buffer) {

		// int bl = Buffer.length;
		byte addrBuf[] = mDGPack.getAddress().getAddress();
		int dgPort = mDGPack.getPort();
		int headerLen = 6 + addrBuf.length;
		int dataLen = mDGPack.getLength();
		int packageLen = headerLen + dataLen;

		byte UB[] = new byte[packageLen];

		UB[0] = (byte) 0x00; // 保留位 0x00
		UB[1] = (byte) 0x00; // 保留位 0x00
		UB[2] = (byte) 0x00; // FRAG '00' - Standalone DataGram
		UB[3] = (byte) 0x01; // 地址类型 -->'01'-IP v4
		System.arraycopy(addrBuf, 0, UB, 4, addrBuf.length);
		UB[4 + addrBuf.length] = (byte) ((dgPort >> 8) & 0xFF);
		UB[5 + addrBuf.length] = (byte) ((dgPort) & 0xFF);
		System.arraycopy(buffer, 0, UB, 6 + addrBuf.length, dataLen);
		System.arraycopy(UB, 0, buffer, 0, packageLen);

		return UB;

	}

	private byte[] clearDgpHead(byte[] buffer) {
		int addressLength = 0;
		int pos = 4;

		byte addressType = buffer[3]; // IP Address Type
		switch (addressType) {
		case 0x01:
			addressLength = 4;
			break;
		case 0x03:
			addressLength = buffer[pos] + 1;
			break;
		default:
			DebugLog.getInstance()
					.println("Error in ClearDGPhead() - Invalid Destination IP Addres type " + addressType);
			return null;
		}

		byte addressBuf[] = new byte[addressLength];
		System.arraycopy(buffer, pos, addressBuf, 0, addressLength);
		pos += addressLength;

		mUDP_IA = calcInetAddress(addressType, addressBuf);
		mUDP_port = Utils.calcPort(buffer[pos++], buffer[pos++]);

		if (mUDP_IA == null) {
			DebugLog.getInstance().println("Error in ClearDGPHead() - Invalid UDP dest IP address: NULL");
			return null;
		}

		int dataLen = mDGPack.getLength();
		dataLen -= pos;
		byte UB[] = new byte[dataLen];
		System.arraycopy(buffer, pos, UB, 0, dataLen);
		System.arraycopy(UB, 0, buffer, 0, dataLen);

		return UB;

	}

	protected void udpSend(DatagramPacket dgp) {

		if (dgp == null){
			return;
		}
		String LogString = dgp.getAddress() + ":" + dgp.getPort() + "> : " + dgp.getLength() + " bytes";
		try {
			mDGSocket.send(dgp);
		} catch (IOException e) {
			DebugLog.getInstance().println("Error in ProcessUDPClient() - Failed to Send DGP to " + LogString);
			return;
		}
	}

	public void processUdp() {

		// 尝试接收udp数据
		try {
			mDGSocket.receive(mDGPack);
		} catch (InterruptedIOException e) {
			return; // Time Out
		} catch (IOException e) {
			DebugLog.getInstance().println("Error in ProcessUDP() - " + e.toString());
			return;
		}

		if (mClientIP.equals(mDGPack.getAddress())) {
			processUdpClient();
		} else {
			processUdpServer();
		}

		try {
			initUdpInOut(); 
		} catch (IOException e) {
			DebugLog.getInstance().println("IOError in Init_UDP_IO() - " + e.toString());
			mProxyHandler.close();
		}
	}

	/**
	 * 处理客户端的udp数据
	 */
	public void processUdpClient() {

		mClientPort = mDGPack.getPort();
		byte[] Buf = clearDgpHead(mDGPack.getData());
		if (Buf == null)
			return;

		if (Buf.length <= 0)
			return;

		if (mUDP_IA == null) {
			DebugLog.getInstance().println("Error in ProcessUDPClient() - Invalid Destination IP - NULL");
			return;
		}
		if (mUDP_port == 0) {
			DebugLog.getInstance().println("Error in ProcessUDPClient() - Invalid Destination Port - 0");
			return;
		}

		if (mServer != mUDP_IA || mServerPort != mUDP_port) {
			mServer = mUDP_IA;
			mServerPort = mUDP_port;
		}

		DebugLog.getInstance().println("Datagram : " + Buf.length + " bytes : " + Utils.getSocketInfo(mDGPack) + " >> <"
				+ Utils.iP2Str(mServer) + ":" + mServerPort + ">");

		DatagramPacket DGPSend = new DatagramPacket(Buf, Buf.length, mUDP_IA, mUDP_port);

		udpSend(DGPSend);
	}

	public void processUdpServer() {

		DebugLog.getInstance()
				.println("Datagram : " + mDGPack.getLength() + " bytes : " + Utils.getSocketInfo(mDGPack));

		InetAddress DGP_IP = mDGPack.getAddress();
		int DGP_Port = mDGPack.getPort();

		byte[] udpBuf  = addDgpHead(mProxyHandler.mBuffer);
		if (udpBuf == null)
			return;

		// SendTo Client
		DatagramPacket _DGP = new DatagramPacket(udpBuf, udpBuf.length, mClientIP, mClientPort);
		udpSend(_DGP);

		if (DGP_IP != mUDP_IA || DGP_Port != mUDP_port) {
			mServer = DGP_IP;
			mServerPort = DGP_Port;
		}
	}
	/**
	 * 响应客户端-成功 返回 0x00
	 */
	public byte getSuccessCode() {
		return 0x00;
	}
	/**
	 * 响应客户端-失败 返回 0x04
	 */
	public byte getFailCode() {
		return 0x04;
	}

}
